package com.xagent.engine.hook.transformer;

import com.xagent.engine.context.AgentContext;
import com.xagent.engine.context.inter.Context;
import com.xagent.engine.hook.HookInfo;
import com.xagent.engine.hook.enhance.v1.CommonV1Adapter;
import org.objectweb.asm.ClassReader;
import org.objectweb.asm.ClassVisitor;
import org.objectweb.asm.ClassWriter;

import java.io.File;
import java.io.FileOutputStream;
import java.io.IOException;
import java.lang.instrument.ClassFileTransformer;
import java.lang.instrument.IllegalClassFormatException;
import java.security.ProtectionDomain;
import java.util.HashMap;

public class DefaultTransformer implements ClassFileTransformer {
    private final Context context = AgentContext.getContext();

    @Override
    public byte[] transform(ClassLoader loader, String className, Class<?> classBeingRedefined, ProtectionDomain protectionDomain, byte[] classfileBuffer) throws IllegalClassFormatException {
        className = className.replaceAll("/", ".");
        HashMap<String, HashMap<String, HookInfo>> hooksMap = context.getConfig().getPolicyMap();

        // 不需要Hook的包和类
        if (context.getInstWrapper().filterOrEnhance(className)) return new byte[0];

        for (String key : hooksMap.keySet()) {
            if (!hooksMap.get(key).containsKey(className)) continue;
            /**
             * test start
             * */
            ClassReader classReader = new ClassReader(classfileBuffer);
            ClassWriter classWriter = new ClassWriter(ClassWriter.COMPUTE_FRAMES);
            ClassVisitor classVisitor = new CommonV1Adapter(context, className, classWriter);
            classReader.accept(classVisitor, ClassReader.EXPAND_FRAMES);
            byte[] bytes = classWriter.toByteArray();

            if (context.isDebug()) {
                String classPath = context.getConfig().getAgentFilePath().getBaseFilePath() + "/class/";
                File classDir = new File(classPath);
                if (!classDir.exists()) classDir.mkdir();
                String classFile = classPath + className + ".class";

                AgentContext.getContext().getAgentLogger().debug("[DEBUG] " + "recording class " + className + " to :" + classFile);
                try {
                    FileOutputStream fileOutputStream = new FileOutputStream(classFile);
                    fileOutputStream.write(bytes);
                    fileOutputStream.close();
                } catch (IOException e) {
                    AgentContext.getContext().getAgentLogger().debug("[DEBUG] " + "recorded class " + className + " failed, exception info:" + e.getMessage());
                }

            }
            if (bytes != null && bytes != new byte[0]) return bytes;
            /**
             * test end
             * */
        }

        return new byte[0];
    }

}
